'''

Module containing Q-learning for the domination game. 

Update function based on Sutton & Barto's Q-Learning:
http://www.cse.iitm.ac.in/~cs670/book/node65.html

'''

from OpenNero import *
import action_script
from constants import *
from custom_functions import *

import os
import cPickle
import math
import random
from operator import itemgetter

class QBrain(AgentBrain):
    '''
    Class that holds the Q-dictionary and the necessary update functions.
    
    Holds reference to the parameters and the pickle file.
    '''
    epsilon = 0
    alpha = 0
    gamma = 0
    filename = ''
    q_space = {}
    action_set = []
    previous_state = None
    fullpath = ''
    current_action = ''
    
    @classmethod
    def loader(cls,filename=None, fullpath=None, epsilon=None, alpha=None, gamma=None):
        '''
        Loads a QSpace instance. Assumes an initialized pickle file in the same
        folder as this file
        
        epsilon, alpha and gamma are only used if no file exists yet and a new
        one is created
        '''
        thispath = os.path.abspath(os.path.dirname(__file__))
        
        if fullpath:
            path = fullpath
        elif filename:
            path = thispath+ os.sep + filename
        else:
            raise Exception("You HAVE to specify either <filename> or <fullpath>")
        
        try:
            f = open(path, 'rb')
        except:
            #File doesn't exist yet
            qspace = cls(epsilon, alpha, gamma, filename, path)
            print "Created new QSpace file at:", qspace.fullpath
        else:
            qspace = cPickle.load(f)
            qspace._set_qfile_pointers(fullpath=path)
            if epsilon:
                qspace.set_epsilon(epsilon)
            if gamma:
                qspace.set_epsilon(gamma)
            if alpha:
                qspace.set_epsilon(alpha)
            f.close()
            print "Loaded QSpace file at:", path
        
        return qspace

    def __init__(self, epsilon, alpha, gamma, qspacefilename=None):#, states, actions, filename=None, fullpath=None):
        '''
        Initializes the Q-Dictionary and sets all parameters.
        '''
        #self.statespace = states
        #self.actionspace = actions
        AgentBrain.__init__(self) # do not remove!
        if qspacefilename is None:
            self.q_space = {}
            self.q_space['epsilon'] = epsilon
            self.q_space['alpha'] = alpha
            self.q_space['gamma'] = gamma
            self._set_qfile_pointers(None)
            self.save()
        else:
            self._set_qfile_pointers(qspacefilename)
            self._load_qspace()
    
    def initialize(self, init_info):
        """
        init_info contains a description of the observations the agent will see
        and the actions the agent will be able to perform
        """
        self.init_info = init_info
        self._init_actionset()
        
        return True
    
    def start(self, time, sensors):
        """
        Take in the initial sensors and return the first action
        """
        action = self.get_action(sensors)
        
        self.previous_action = action
        self.previous_sensors = sensors
        
        return action
    
    def act(self, time, sensors, reward):
        """
        Take in new sensors and reward from previous action and return the next action
        Specifically, just move toward the closest crumb!
        """
        self.update(sensors, reward)
        action = self.get_action(sensors)
        
        self.previous_action = action
        self.previous_sensors = sensors
        
        return action
    
    def end(self, time, reward):
        """
        take in final reward
        """
        self.update(None, reward)
        return True
    
    def destroy(self):
        """
        called when the agent is destroyed
        """
        return True
    
    
    def _set_qfile_pointers(self, filename=None):
        '''
        Function that correctly sets up the path and name of the file in which 
        the instance will be saved
        '''
        
        this_path = os.path.abspath(os.path.dirname(__file__))
        
        if filename:
            self.fullpath = this_path + os.sep + filename
        else:
            self.fullpath = this_path + os.sep + 'QSpace' + \
                '-e' + str(self.q_space['epsilon']).replace('.','_') + \
                'a' + str(self.q_space['alpha']).replace('.','_') + 'g' +\
                str(self.q_space['gamma']).replace('.', '_') + '.pickle'
    
    def _init_actionset(self):
        '''
        Create actionset
        
        Separate function for easy modification
        '''
        mina = self.init_info.actions.min(0)
        maxa = self.init_info.actions.max(0)
        self.action_set = range(int(mina), int(maxa)+1)
    
    def _get_actionset_from_qspace(self):
        '''
        Extracts set of actions from loaded file
        '''
        self.action_set = self.q_space[self.q_space.keys()[0]].keys()
    
    def _init_state(self, state):
        '''
        Create initial q_space dictionary
        '''
        self.q_space[state] = {}
        for a in self.action_set:
            self.q_space[state][a] = 0.0
    
    def _load_qspace(self):
        '''
        Get q_space from file, if it exists.
        '''
        try:
            f = open(self.fullpath, 'rb')
        except:
            print "Use New QSpace file at:", self.fullpath
        else:
            self.q_space = cPickle.load(f)
            f.close()
            print "Loaded QSpace file at:", path
            
    
    def _get_a_v_max(self, state):
        '''
        Returns action-value that has highest q-value in <state>
        '''
        s = tuple(state)
        if s not in self.q_space:
            self._init_state(s)

        return sorted(self.q_space[s].items(), key=itemgetter(1)).pop()
    
    def set_epsilon(self, epsilon):
        '''
        setter
        '''
        self.q_space['epsilon'] = epsilon
    
    def set_gamma(self, gamma):
        '''
        setter
        '''
        self.q_space['gamma'] = gamma
    
    def set_alpha(self, alpha):
        '''
        setter
        '''
        self.q_space['alpha'] = alpha
    
    def get_action(self, sensors):
        '''
        Return the action for the agent
        '''
        action = self.init_info.actions.get_instance()
        # just tell the agent which way to go!
        
        if random.random() < self.q_space['epsilon']:
            nr_actions = len(self.action_set)
            action[0] = self.action_set[int(math.floor(nr_actions*random.random()))]
        else:
            action[0] = self._get_a_v_max(sensors)[0]
        
        return action
    
    def update(self, sensors, reward):
        '''
        Update current state-action Q value Q(s,a) and set previous state
        '''
        p_action = self.previous_action[0]
        p_sensors = tuple(self.previous_sensors)
        
        if p_sensors not in self.q_space:
            self._init_state(p_sensors)
        q_sa = self.q_space[p_sensors][p_action]
        
        if sensors is not None:
            q_max = self._get_a_v_max(sensors)[1]
        else:
            q_max = 0
        self.q_space[p_sensors][p_action] = q_sa + self.q_space['alpha']* (reward[0] + self.q_space['gamma'] * q_max - q_sa ) 
    
    def save(self, verbose=True):
        '''
        Saves pickle with qvalues
        '''
        
        mode = 'wb'
        qspacefile = open(self.fullpath, mode)
        cPickle.dump(self.q_space, qspacefile, cPickle.HIGHEST_PROTOCOL)
        qspacefile.close()
        if verbose:
            print "Instance saved to", self.fullpath
        
